from nets.swin_seg import swin_seg_model
from utils.datasets import MySeg2dDatasets
from utils.learning_rate_set import lr_scheduler
from tensorflow.keras import losses,optimizers
from tensorflow.keras.callbacks import (ModelCheckpoint,TensorBoard,
                                LearningRateScheduler,CSVLogger)
import time
from utils.my_callbacks import DisplayCallback,SaveModelCallback
import os
import numpy as np
from utils.losses import (dice_coefficient,dice_loss)


if __name__ == '__main__':
    input_shape = [224,224,3]
    num_classes = 2
    batch_size = 4
    
    model = swin_seg_model(input_shape,num_classes)
    model.summary()
    # input('zz')
    # loss = losses.SparseCategoricalCrossentropy(from_logits=True)
    model.compile(
        optimizer=optimizers.Adam(learning_rate=1e-3),
        loss=dice_loss,
        metrics=[dice_coefficient],
    )
    

    # ------------------------------------
    # datasets
    # ------------------------------------
    work_dir = "C:/Users/hblee/Documents/datasets/seg2d"
    train_datasets = MySeg2dDatasets(work_dir=work_dir,
                                    batch_size=batch_size,
                                    input_shape=input_shape,
                                    mode='train',
                                    num_classes=num_classes,)

    val_datasets = MySeg2dDatasets(work_dir=work_dir,
                                    batch_size=batch_size,
                                    input_shape=input_shape,
                                    mode='val',
                                    num_classes=num_classes,)
    # -----------------------------------------------------
    # 各种回调函数
    # -----------------------------------------------------
    # tensorboard回调
    tb_callback = TensorBoard(log_dir='logs')

    # 学习率下降回调
    lr_callback = LearningRateScheduler(lr_scheduler)

    # 保存训练过程回调
    t = time.time()
    now = time.localtime(t)
    prefix = f'{now.tm_year:04}{now.tm_mon:02}{now.tm_mday:02}{now.tm_hour:02}{now.tm_min:02}{now.tm_sec:02}'
    file_name = f"logs/log_{prefix}.csv"
    history_logger = CSVLogger(file_name,separator=',',append=True)

    # 模型保存回调
    checkpoint_path = "logs/cp{epoch:03d}_loss{loss:.3f}-valoss{val_loss:.3f}.ckpt"
    checkpoint_dir = os.path.dirname(checkpoint_path)
    cp_callback = ModelCheckpoint(filepath=checkpoint_path,
                                save_weights_only=True,
                                save_best_only=True,
                                verbose=1)


    # --------------------------------------------
    # 训练
    # --------------------------------------------
    history = model.fit(train_datasets, 
                        validation_data=val_datasets,
                        epochs=10,
                        callbacks=[ tb_callback,
                                    history_logger,
                                    DisplayCallback(model),
                                    cp_callback,
                                    # SaveModelCallback(model),
                                    ])
    # model.save("model.h5")
    np.save('res.npy',history.history)
    print('训练结束'.center(60,'='))

